#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
下载YOLO模型到MaixCAM设备的脚本
"""

import os
import requests
from maix import fs

# YOLO模型下载链接（来自MaixHub模型库）
MODEL_URLS = {
    "yolov5s.mud": "https://maixhub.com/model/download/365/1/1",
    "yolov8n.mud": "https://maixhub.com/model/download/400/1/1",
    "yolo11n.mud": "https://maixhub.com/model/download/453/1/1"
}

# 模型保存路径
MODEL_DIR = "/root/models/"
DEFAULT_MODEL = "yolov5s.mud"

def create_model_dir():
    """创建模型目录"""
    try:
        if not os.path.exists(MODEL_DIR):
            os.makedirs(MODEL_DIR)
            print(f"创建模型目录: {MODEL_DIR}")
        return True
    except Exception as e:
        print(f"创建模型目录失败: {e}")
        return False

def download_model(model_name, url):
    """下载指定的模型文件"""
    try:
        print(f"正在下载 {model_name}...")
        
        # 发送HTTP请求
        response = requests.get(url, timeout=30)
        response.raise_for_status()  # 检查请求是否成功
        
        # 保存文件
        model_path = os.path.join(MODEL_DIR, model_name)
        with open(model_path, 'wb') as f:
            f.write(response.content)
        
        print(f"{model_name} 下载完成，保存到: {model_path}")
        return True
        
    except requests.exceptions.RequestException as e:
        print(f"下载 {model_name} 失败: {e}")
        return False
    except Exception as e:
        print(f"保存 {model_name} 时出错: {e}")
        return False

def check_model_exists(model_name):
    """检查模型文件是否已存在"""
    model_path = os.path.join(MODEL_DIR, model_name)
    return os.path.exists(model_path)

def main():
    """主函数"""
    print("MaixCAM YOLO模型下载工具")
    print("=" * 30)
    
    # 创建模型目录
    if not create_model_dir():
        print("无法创建模型目录，程序退出")
        return
    
    # 下载默认模型（YOLOv5s）
    model_name = DEFAULT_MODEL
    if not check_model_exists(model_name):
        if download_model(model_name, MODEL_URLS[model_name]):
            print(f"{model_name} 准备就绪")
        else:
            print(f"{model_name} 下载失败")
    else:
        print(f"{model_name} 已存在，跳过下载")
    
    print("模型下载完成！")
    print("现在可以运行 yolo_camera.py 进行对象检测了")

if __name__ == "__main__":
    main()